# -*- coding: utf-8 -*-
"""DenseNet models for TF Keras.

# Reference paper

- [Densely Connected Convolutional Networks]
  (https://arxiv.org/abs/1608.06993) (CVPR 2017 Best Paper Award)

# Reference implementation

- [Torch DenseNets]
  (https://github.com/liuzhuang13/DenseNet/blob/master/models/densenet.lua)
- [TensorNets]
  (https://github.com/taehoonlee/tensornets/blob/master/tensornets/densenets.py)
"""
import tensorflow as tf
from tensorflow import keras

Input = keras.layers.Input
ZeroPadding2D = keras.layers.ZeroPadding2D
Activation = keras.layers.Activation
MaxPooling2D = keras.layers.MaxPooling2D
add = keras.layers.add
Dense = keras.layers.Dense
Conv2D = keras.layers.Conv2D
BatchNormalization = keras.layers.BatchNormalization
Flatten = keras.layers.Flatten
AveragePooling2D = keras.layers.AveragePooling2D
Concatenate = keras.layers.Concatenate
Model = keras.models.Model


def dense_block(x, blocks, name):
    """A dense block.

    # Arguments
        x: input tensor.
        blocks: integer, the number of building blocks.
        name: string, block label.

    # Returns
        output tensor for the block.
    """
    for i in range(blocks):
        x = conv_block(x, 32, name=name + '_block' + str(i + 1))
    return x


def transition_block(x, reduction, name):
    """A transition block.

    # Arguments
        x: input tensor.
        reduction: float, compression rate at transition layers.
        name: string, block label.

    # Returns
        output tensor for the block.
    """
    bn_axis = 3
    x = BatchNormalization(axis=bn_axis, epsilon=1.001e-5,
                           name=name + '_bn')(x)
    x = Activation('relu', name=name + '_relu')(x)
    x = Conv2D(int(x.shape.as_list()[bn_axis] * reduction), 1,
               use_bias=False, name=name + '_conv')(x)
    x = AveragePooling2D(2, strides=2, name=name + '_pool')(x)
    return x


def conv_block(x, growth_rate, name):
    """A building block for a dense block.

    # Arguments
        x: input tensor.
        growth_rate: float, growth rate at dense layers.
        name: string, block label.

    # Returns
        Output tensor for the block.
    """
    bn_axis = 3
    x1 = BatchNormalization(axis=bn_axis, epsilon=1.001e-5,
                            name=name + '_0_bn')(x)
    x1 = Activation('relu', name=name + '_0_relu')(x1)
    x1 = Conv2D(4 * growth_rate, 1, use_bias=False,
                name=name + '_1_conv')(x1)
    x1 = BatchNormalization(axis=bn_axis, epsilon=1.001e-5,
                            name=name + '_1_bn')(x1)
    x1 = Activation('relu', name=name + '_1_relu')(x1)
    x1 = Conv2D(growth_rate, 3, padding='same', use_bias=False,
                name=name + '_2_conv')(x1)
    x = Concatenate(axis=bn_axis, name=name + '_concat')([x, x1])
    return x


def densenet(blocks, input_shape=(224, 224, 3), output_shape=10, bn_axis=3):
    """
    :param blocks:
        [6, 12, 24, 16]:  densenet121
        [6, 12, 32, 32]:  densenet169
        [6, 12, 48, 32]:  densenet201

    :param input_shape:
        图片大小

    :param output_shape:
        图片类别数量

    :param bn_axis:
        bn_axis = 3 if backend.image_data_format() == 'channels_last' else 1

    :return:
    """

    img_input = Input(shape=input_shape)

    with tf.name_scope('DenseNet'):
        x = ZeroPadding2D(padding=((3, 3), (3, 3)))(img_input)
        x = Conv2D(64, 7, strides=2, use_bias=False, name='conv1/conv')(x)
        x = BatchNormalization(
            axis=bn_axis, epsilon=1.001e-5, name='conv1/bn')(x)
        x = Activation('relu', name='conv1/relu')(x)
        x = ZeroPadding2D(padding=((1, 1), (1, 1)))(x)
        x = MaxPooling2D(3, strides=2, name='pool1')(x)

        with tf.name_scope('dense_block_1'):
            x = dense_block(x, blocks[0], name='conv2')

        with tf.name_scope('transition_block_1'):
            x = transition_block(x, 0.5, name='pool2')

        with tf.name_scope('dense_block_2'):
            x = dense_block(x, blocks[1], name='conv3')

        with tf.name_scope('transition_block_2'):
            x = transition_block(x, 0.5, name='pool3')

        with tf.name_scope('dense_block_3'):
            x = dense_block(x, blocks[2], name='conv4')

        with tf.name_scope('transition_block_3'):
            x = transition_block(x, 0.5, name='pool4')

        with tf.name_scope('dense_block_4'):
            x = dense_block(x, blocks[3], name='conv5')

        x = BatchNormalization(
            axis=bn_axis, epsilon=1.001e-5, name='bn')(x)
        x = Activation('relu', name='relu')(x)

    # pooling = None
    # if pooling == 'avg':
    #     x = layers.GlobalAveragePooling2D(name='avg_pool')(x)
    # elif pooling == 'max':
    #     x = layers.GlobalMaxPooling2D(name='max_pool')(x)

    x = AveragePooling2D((7, 7), name='avg_pool')(x)
    x = Flatten()(x)
    predictions = Dense(output_shape, activation='softmax', name='my_fc2')(x)

    # Create model.
    if blocks == [6, 12, 24, 16]:
        model = Model(img_input, predictions, name='densenet121')
    elif blocks == [6, 12, 32, 32]:
        model = Model(img_input, predictions, name='densenet169')
    elif blocks == [6, 12, 48, 32]:
        model = Model(img_input, predictions, name='densenet201')
    else:
        model = Model(img_input, predictions, name='densenet')

    return model


# 创建模型
def inference(input_shape, output_shape, trainable=False, drop_out=None):
    """
    :param input_shape:   输入数据大小
    :param output_shape:  输出数据大小
    :param trainable      参数是否可重新训练
    :param drop_out       Drop Out 比例， 为 None 时则不使用 DropOut
    :return:
    """
    del drop_out
    del trainable

    densenet121 = densenet(blocks=[6, 12, 24, 16],
                           input_shape=(input_shape, input_shape, 3),
                           output_shape=output_shape)
    return densenet121

if __name__ == '__main__':
    _model = inference(224, 10, trainable=False, drop_out=None)
    # 打印模型结构
    _model.summary()
    import json
    print('\nJSON Model Config:\n',
          json.dumps(json.loads(_model.to_json()), ensure_ascii=False, sort_keys=True, indent=4))
